from __future__ import print_function

import os
import time
import tensorflow as tf
from flask import Flask, render_template, request, send_from_directory

import model
import reader
from preprocessing import preprocessing_factory

app = Flask(__name__)
app.config['SECRET_KEY'] = '123456'
app.static_folder = 'static'

UPLOAD_FOLDER = 'static/img/uploads/'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

tf.app.flags.DEFINE_string('loss_model', 'vgg_16', 'The name of the architecture to evaluate. '
                                                   'You can view all the support models in nets/nets_factory.py')
tf.app.flags.DEFINE_integer('image_size', 256, 'Image size to train.')
tf.app.flags.DEFINE_string("model_file", "models.ckpt", "")
tf.app.flags.DEFINE_string("image_file", "a.jpg", "")
FLAGS = tf.app.flags.FLAGS


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS


@app.route('/')#路由修饰器，触发index.html
def index():
    return render_template('index.html')#渲染


@app.route('/transform', methods=['GET', 'POST'])#动态URL
def deal_image():
    # model 字典
    models_dict = {'cubist': 'cubist.ckpt-done',
                   'denoised_starry': 'denoised_starry.ckpt-done',
                   'feathers': 'feathers.ckpt-done',
                   'mosaic': 'mosaic.ckpt-done',
                   'scream': 'scream.ckpt-done',
                   'udnie': 'udnie.ckpt-done',
                   'wave': 'wave.ckpt-done',
                   'painting': 'painting.ckpt-done',
                   'Wheat': 'Wheat.ckpt-3000',
                   }
    if request.method == 'POST':
        # 响应上传图片pic
        file = request.files['pic']#上传图片
        # 响应选择的style
        style = request.form['style']#表单id
        # 判断文件合法性
        if file and allowed_file(file.filename):
            # 检查上传路径是否存在，不存在创建
            if os.path.exists(app.config['UPLOAD_FOLDER']) is False:
                os.makedirs(app.config['UPLOAD_FOLDER'])
            file.save(os.path.join(app.config['UPLOAD_FOLDER'], file.filename))#将上传的文件保存到目录
            # 初始化model_file
            model_file = 'wave.ckpt-done'
            if style != '':
                if models_dict[style] != '':#若获取的风格不为空则赋值
                    model_file = models_dict[style]
            # 风格转换
            style_transform(style, 'models/' + model_file, os.path.join(app.config['UPLOAD_FOLDER']) + file.filename,
                            style + '_res_' + file.filename)
            # 转换完后渲染到transformed.html页面，显示三张图片
            return render_template('transformed.html', style='img/style/' + style + '.jpg',
                                   upload='img/uploads/' + file.filename,
                                   transformed='img/generated/' + style + '_res_' + file.filename)
        return 'transform error:file format error'
    return 'transform error:method not post'


@app.route('/uploads/<filename>')
#上传转换好的图片
def uploaded_file(filename):#主要用于下载文件
    return send_from_directory('static/img/generated/', filename)

#风格转换函数style：风格, model_file：风格模型, img_file：输入图片, result_file：结果图片
def style_transform(style, model_file, img_file, result_file):
    height = 0
    width = 0
    with open(img_file, 'rb') as img:
        with tf.Session().as_default() as sess:
            if img_file.lower().endswith('png'):
                image = sess.run(tf.image.decode_png(img.read()))
            else:
                image = sess.run(tf.image.decode_jpeg(img.read()))
            height = image.shape[0]
            width = image.shape[1]
    print('Image size: %dx%d' % (width, height))

    with tf.Graph().as_default():
        with tf.Session().as_default() as sess:
            image_preprocessing_fn, _ = preprocessing_factory.get_preprocessing(
                FLAGS.loss_model,
                is_training=False)
            image = reader.get_image(img_file, height, width, image_preprocessing_fn)
            image = tf.expand_dims(image, 0)#增加一个维度
            generated = model.net(image, training=False)
            generated = tf.squeeze(generated, [0])#删除一个维度
            saver = tf.train.Saver(tf.global_variables())#初始化一个saver
            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
            FLAGS.model_file = os.path.abspath(model_file)
            saver.restore(sess, FLAGS.model_file)
            '''开始转换'''
            start_time = time.time()
            generated = sess.run(generated)
            generated = tf.cast(generated, tf.uint8)
            end_time = time.time()
            print('Elapsed time: %fs' % (end_time - start_time))
            '''生成图片存放位置'''
            generated_file = 'static/img/generated/' + result_file
            if os.path.exists('static/img/generated') is False:
                os.makedirs('static/img/generated')
            with open(generated_file, 'wb') as img:
                img.write(sess.run(tf.image.encode_jpeg(generated)))
                print('Done. Please check %s.' % generated_file)


if __name__ == '__main__':
    app.run(debug=True)



